import torch
from torchvision import transforms as T
from torch.utils.data import Dataset, DataLoader


class MyMNIST(Dataset):
    def __init__(self, images, labels, transformer=None):
        super().__init__()
        self.images = images
        self.labels = labels
        self.transformer = transformer

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        img = self.images[idx]
        label = self.labels[idx]
        if self.transformer:
            img = self.transformer(img)
        return img, label


# 准备数据
images = torch.rand(100, 1, 32, 32).float()
labels = torch.randint(0, 10, (100,))
transformer = T.Compose([
    T.Resize((28, 28), antialias=True)
])

# 实例化对象
ds = MyMNIST(images, labels, transformer)

# 查看数据集的长度
print(len(ds))

# 查看数据集的元素
img, label = ds[0:10]
print(f"图像的大小 {img.shape} 对应的标签：{label.shape}")

# 1. 学生练习：使用random_split 按照8:2划分训练集和验证集
train_ds, val_ds = torch.utils.data.random_split(ds, [round(len(ds)*0.8), round(len(ds)*0.2)])

# 2. 学生练习：使用DataLoader，创建训练集 和 验证集的加载器，并验证加载器是否成功
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=32, shuffle=True)
for img, label in train_loader:
    print(f"图像的大小 {img.shape} 对应的标签：{label.shape}")
    break
for img, label in val_loader:
    print(f"图像的大小 {img.shape} 对应的标签：{label.shape}")
    break
